# central settings
#
simulation     = True
year           = "2012"
magnetPolarity = "Up"
###


########################################################################
from os import environ
import GaudiKernel.SystemOfUnits as Units
from Gaudi.Configuration import *
from Configurables import (GaudiSequencer,
                           DecayTreeTuple,
                           CombineParticles,
                           LoKi__Hybrid__TupleTool,
                           DaVinci,
                           FilterDesktop)

from Configurables import LoKi__Hybrid__PlotTool as PlotTool
from Configurables import LoKi__Hybrid__FilterCriterion as LoKiFilterCriterion
from DecayTreeTuple.Configuration import *

from GaudiConfUtils.ConfigurableGenerators import  CombineParticles
from PhysSelPython.Wrappers import (SelectionSequence, Selection, DataOnDemand, AutomaticData, MergedSelection)
from Configurables import LoKi__VoidFilter as Filter
from PhysConf.Filters import LoKi_Filters
LoKiTool = LoKi__Hybrid__TupleTool( 'LoKiTool')
#
######################################################################################################

if simulation is True:
    location = 'AllStreams/Phys/FullDSTDiMuonJpsi2MuMuDetachedLine/Particles'#MC的TES位置
if simulation is False:
    location = 'Dimuon/Phys/FullDSTDiMuonJpsi2MuMuDetachedLine/Particles'#data的TES位置
######################################################################################################
# if you are setting something wrong, it should die here...
if year not in [ "2011", "2012" ]:
    raise RuntimeError(year + " is not a valid year")
if magnetPolarity not in ["Up", "Down"]:
    raise RuntimeError(magnetPolarity + " is not a valid magnet polarity")
######################################################################################################
# CUTS CUTS CUTS CUTS CUTS CUTS CUTS
####################################
BdCombCuts = "(AM > 4500 * MeV) & "\
            "(AM < 7100 * MeV)"
BdCuts = "(M                                   > 4600 * MeV) & "\
    "(M                                   < 7000 * MeV) & "\
    "(BPVIPCHI2()                         < 16) & "\
    "(BPVDIRA                             > 0.9995) & "\
    "(BPVVDCHI2                           > 64) & "\
    "(MAXTREE(ISBASIC,MIPCHI2DV(PRIMARY)) > 9 ) & "\
    "(VFASPF(VCHI2/VDOF)                  < 10)"

MuonCuts = "(TRGHP < 0.5) & "\
          "(MIPCHI2DV(PRIMARY) > 9) & "\
          "(PIDmu> 0) & "\
          "(ISMUON)"

Pi0Cuts = "(PT > 300) &"\
    "(CL > 0.02)"

GammaCuts = "(PT > 200) "
######################################################################################################
###
def __Gamma__(GammaCuts):
    _gamma = AutomaticData(Location = 'Phys/StdLooseAllPhotons/Particles')
    _filter = FilterDesktop("GammaFilter", Code = GammaCuts)
    _sel = Selection("Selection_StdLooseAllPhotons",
                     RequiredSelections = [ _gamma ] ,
                     Algorithm = _filter)                    
    return _sel

Gamma = __Gamma__(GammaCuts)

def __Pi0merged__(Gamma, Pi0Cuts):
    _pi0merged = CombineParticles()
    _pi0merged.DecayDescriptor = "[pi0 ->  gamma  gamma]cc"
    _pi0merged.MotherCut = Pi0Cuts
    _pi0Conf = _pi0merged.configurable("Combine_GammaGamma")
    _selpi0merged = Selection("Selection_pi0merged",
                                RequiredSelections = [Gamma],
                                Algorithm = _pi0Conf)
    return _selpi0merged

def __Pi0__(Pi0Cuts):
    _pi0resolved = AutomaticData(Location = 'Phys/StdLooseResolvedPi0/Particles')
    _filter_pi0resolved = FilterDesktop("ResolvedPi0Filter", Code = Pi0Cuts )
    _selpi0resolved = Selection("Selection_pi0resolved",
                                RequiredSelections = [ _pi0resolved ] ,
                                Algorithm = _filter_pi0resolved) 
    _selpi0merged = __Pi0merged__(Gamma, Pi0Cuts)   
    _sel = MergedSelection("Selection_pi0",
                            RequiredSelections = [ _selpi0resolved, _selpi0merged ])
    return _sel

def __jpsi__(location):
    return AutomaticData(Location = location)

###
def __jpsipi0__(Jpsi,Pi0, BdCombCuts, BdCuts):
    _jpsipi0 = CombineParticles()
    _jpsipi0.DecayDescriptor = "[B0 -> J/psi(1S) pi0]cc"
    _jpsipi0.CombinationCut = BdCombCuts
    _jpsipi0.MotherCut = BdCuts
    _jpsipi0Conf = _jpsipi0.configurable("Combine_BJpsipi0")
    _selJpsipi0 = Selection("BJpsipi0",
                              Algorithm = _jpsipi0Conf,
                              RequiredSelections = [ Jpsi, Pi0])
    return _selJpsipi0

#######
Pi0             = __Pi0__(Pi0Cuts)
Jpsi            = __jpsi__(location)
Jpsipi0   = __jpsipi0__(Jpsi,Pi0,BdCombCuts, BdCuts)
BJpsipi0Seq = SelectionSequence( 'BJpsipi0Seq', TopSelection = Jpsipi0 )

######################################################################################################
BJpsiPi0 = DecayTreeTuple("BJpsiPi0Tuple")
BJpsiPi0.Inputs = [ Jpsipi0.outputLocation() ]
BJpsiPi0.Decay = "[B0 -> ^(J/psi(1S) -> ^mu+ ^mu-) ^(pi0 -> ^gamma ^gamma)]CC"

# define the tools and configure them
BJpsiPi0.ToolList = [
    "TupleToolKinematic"
    ,"TupleToolPid"
    ,"TupleToolANNPID"
    ,"TupleToolTrackInfo"
    ,"TupleToolRecoStats"
    ,"TupleToolTrigger"
    ,"TupleToolPrimaries"
    ]

# define the list of triggers that could have fired...
triggerList = [ "L0HadronDecision",
                "L0DiMuonDecision",
                "L0MuonDecision",
                "Hlt1TrackMuonDecision",
                "Hlt1TrackAllL0Decision",
                "Hlt1DiMuonLowMassDecision",
                "Hlt1DiMuonHighMassDecision",
                "Hlt1SingleMuonHighPTDecision",
                "Hlt2TopoMu2BodyBBDTDecision",
                "Hlt2TopoMu3BodyBBDTDecision",
                "Hlt2TopoMu4BodyBBDTDecision",
                "Hlt2TopoE2BodyBBDTDecision",
                "Hlt2TopoE3BodyBBDTDecision",
                "Hlt2TopoE4BodyBBDTDecision",
                "Hlt2DiMuonDetachedHeavyDecision",
                "Hlt2SingleMuonDecision",
                "Hlt2DiMuonDetachedDecision",
                "Hlt2DiMuonDetachedJPsiDecision",
                "Hlt2DiMuonDetachedHeavyDecision" ]

# other variables
# including isolation variables
# for microDST, they are stored under a special location
BJpsiPi0.addTupleTool("LoKi::Hybrid::TupleTool/LoKiTool")
BJpsiPi0.LoKiTool.Variables = {
    "eta"                     : "ETA",
    "phi"                     : "PHI",
    "LV01"                    : "LV01",
    "BPVDLS"                  : "BPVDLS" }


# some truthmatching stuff
if simulation is True:
    MCTruth1=BJpsiPi0.addTupleTool("TupleToolMCTruth/MCTruth1")
    MCTruth1.addTupleTool("MCTupleToolHierarchy")


####

####

################################################################
################################################################
# define the branches for resolved pi0
# for resolved pi0
BJpsiPi0.addBranches({  # remove all "^" except where needed.
        "B"          : "^[B0 -> (J/psi(1S) ->  mu+  mu-) ( pi0 ->  gamma  gamma)]CC",
        "mu1"     : " [B0 -> (J/psi(1S) ->  mu+ ^mu-) ( pi0 ->  gamma  gamma)]CC",
        "mu2"    : " [B0 -> (J/psi(1S) -> ^mu+  mu-) ( pi0 ->  gamma  gamma)]CC",
        "Jpsi"       : " [B0 -> ^(J/psi(1S) -> mu+  mu-) ( pi0 ->  gamma  gamma)]CC",
        "pi0": " [B0 -> (J/psi(1S) ->  mu+  mu-) ^(pi0 ->  gamma  gamma)]CC",
        "gamma0"     : " [B0 -> (J/psi(1S) ->  mu+  mu-) ( pi0 -> ^gamma  gamma)]CC",
        "gamma1"     : " [B0 -> (J/psi(1S) ->  mu+  mu-) ( pi0 ->  gamma ^gamma)]CC"
        })

### B tools
BJpsiPi0.addTupleTool("TupleToolDecayTreeFitter/BDTF")
BJpsiPi0.BDTF.constrainToOriginVertex = True
BJpsiPi0.BDTF.daughtersToConstrain = [ 'J/psi(1S)', 'pi0' ]
BJpsiPi0.addTupleTool("TupleToolGeometry")
BJpsiPi0.addTupleTool("TupleToolPropertime")
BJpsiPi0.addTupleTool( "TupleToolTISTOS/TupleToolTISTOS")
BJpsiPi0.TupleToolTISTOS.TriggerList = triggerList
BJpsiPi0.TupleToolTISTOS.VerboseL0 = True
BJpsiPi0.TupleToolTISTOS.VerboseHlt1 = True
BJpsiPi0.TupleToolTISTOS.VerboseHlt2 = True
BJpsiPi0.addTupleTool("TupleToolVtxIsoln")
BJpsiPi0.addTupleTool("TupleToolConeIsolation/TupleToolConeIsolation")
BJpsiPi0.TupleToolConeIsolation.Verbose = True
BJpsiPi0.TupleToolConeIsolation.MinConeSize = 0.6
BJpsiPi0.TupleToolConeIsolation.MaxConeSize = 2.0
BJpsiPi0.TupleToolConeIsolation.SizeStep = 0.3


### Jpsi tools
BJpsiPi0.Jpsi.addTupleTool("TupleToolGeometry")
BJpsiPi0.Jpsi.addTupleTool( "TupleToolTISTOS/TupleToolTISTOS")
BJpsiPi0.Jpsi.TupleToolTISTOS.TriggerList = triggerList
BJpsiPi0.Jpsi.TupleToolTISTOS.VerboseL0 = True
BJpsiPi0.Jpsi.TupleToolTISTOS.VerboseHlt1 = True
BJpsiPi0.Jpsi.TupleToolTISTOS.VerboseHlt2 = True

# mu1tools
BJpsiPi0.mu1.addTupleTool("TupleToolGeometry")
BJpsiPi0.mu1.addTupleTool( "TupleToolTISTOS/TupleToolTISTOS")
BJpsiPi0.mu1.TupleToolTISTOS.TriggerList = triggerList
BJpsiPi0.mu1.TupleToolTISTOS.VerboseL0 = True
BJpsiPi0.mu1.TupleToolTISTOS.VerboseHlt1 = True
BJpsiPi0.mu1.TupleToolTISTOS.VerboseHlt2 = True

# mu2 tools
BJpsiPi0.mu2.addTupleTool("TupleToolGeometry")
BJpsiPi0.mu2.addTupleTool( "TupleToolTISTOS/TupleToolTISTOS")
BJpsiPi0.mu2.TupleToolTISTOS.TriggerList = triggerList
BJpsiPi0.mu2.TupleToolTISTOS.VerboseL0 = True
BJpsiPi0.mu2.TupleToolTISTOS.VerboseHlt1 = True
BJpsiPi0.mu2.TupleToolTISTOS.VerboseHlt2 = True

# pi0 tools
BJpsiPi0.pi0.addTupleTool("TupleToolPi0Info")

# gamma tools
BJpsiPi0.gamma0.addTupleTool("TupleToolPhotonInfo")
BJpsiPi0.gamma0.addTupleTool("TupleToolCaloHypo")

# gamma tools
BJpsiPi0.gamma1.addTupleTool("TupleToolPhotonInfo")
BJpsiPi0.gamma1.addTupleTool("TupleToolCaloHypo")

################################################################
################################################################
Seq1 = GaudiSequencer("Seq1")
Seq1.Members.append( BJpsipi0Seq.sequence() )
Seq1.Members.append( BJpsiPi0 )
Seq1.ModeOR = True
Seq1.ShortCircuit = False
# ######################################################################################################
DaVinci().MainOptions = ""
DaVinci().UserAlgorithms = [Seq1]
# ######################################################################################################
stripFilter = LoKi_Filters (
    STRIP_Code = " HLT_PASS_RE ('StrippingFullDSTDiMuonJpsi2MuMuDetachedLineDecision') "
)
stripFilterSeq = stripFilter.sequence( "StripFilter" )
########################################################################
DaVinci().EvtMax    = -1
DaVinci().DataType  = year
DaVinci().PrintFreq = 1000
#DaVinci().Lumi      = True
DaVinci().EventPreFilters = [ stripFilterSeq ]
DaVinci().InputType = "DST"

if simulation is True:
    DaVinci().Simulation = True
    DaVinci().Lumi       = False
    DaVinci().TupleFile  = "Tuple.root"

if simulation is False:
    DaVinci().TupleFile = "Tuple.root"

# db tags                                                                                                                                             
if simulation is True:
    if year is "2011":
        if magnetPolarity is "Down":
            DaVinci().CondDBtag = "sim-20130522-vc-md100"
            DaVinci().DDDBtag   = "dddb-20130929"
        if magnetPolarity is "Up":
            DaVinci().CondDBtag = "sim-20130522-vc-mu100"
            DaVinci().DDDBtag   = "dddb-20130929"

    if year is "2012":
        if magnetPolarity is "Down":
            DaVinci().CondDBtag = "sim-20160321-2-vc-md100"
            DaVinci().DDDBtag   = "dddb-20170721-2"
        if magnetPolarity is "Up":
            DaVinci().CondDBtag = "sim-20160321-2-vc-mu100"
            DaVinci().DDDBtag   = "dddb-20170721-2"

### mag polarity does not matter                                                                                                                       
if simulation is False:
    if year is "2011":
        DaVinci().CondDBtag = "cond-20141107"
        DaVinci().DDDBtag   = "dddb-20130929"

    if year is "2012":
        DaVinci().CondDBtag = "cond-20141107"
        DaVinci().DDDBtag   = "dddb-20130929-1"

